#!/usr/bin/env python3
"""
测试数据生成脚本
用于生成测试数据以验证系统功能
"""

import sys
import os
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, project_root)

import asyncio
import random
from datetime import datetime, timedelta
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy import select
from sqlmodel import SQLModel
from database.models import (
    User, Product, Order, Manufacturer, Model3D, ProductionJob, 
    ProductImage, Category, UserRole, OrderStatus, JobStatus
)
from database.connection import DATABASE_URL
import bcrypt


async def create_test_users(session: AsyncSession):
    """创建测试用户"""
    print("创建测试用户...")
    
    # 检查是否已存在用户
    result = await session.execute(select(User))
    existing_users = result.scalars().all()
    
    if len(existing_users) > 1:  # 除了管理员外还有其他用户
        print("测试用户已存在，跳过创建")
        return
    
    # 创建测试用户
    users = []
    
    # 普通用户
    for i in range(1, 6):
        hashed_password = bcrypt.hashpw(f"user{i}123".encode('utf-8'), bcrypt.gensalt())
        user = User(
            email=f"user{i}@example.com",
            username=f"user{i}",
            hashed_password=hashed_password.decode('utf-8'),
            full_name=f"测试用户{i}",
            role=UserRole.CUSTOMER,
            is_active=True,
            phone=f"1380000000{i}",
            address=f"测试地址{i}号"
        )
        users.append(user)
    
    # 厂商用户
    for i in range(1, 4):
        hashed_password = bcrypt.hashpw(f"manufacturer{i}123".encode('utf-8'), bcrypt.gensalt())
        user = User(
            email=f"manufacturer{i}@example.com",
            username=f"manufacturer{i}",
            hashed_password=hashed_password.decode('utf-8'),
            full_name=f"厂商用户{i}",
            role=UserRole.MANUFACTURER,
            is_active=True,
            phone=f"1390000000{i}",
            company_name=f"测试厂商{i}"
        )
        users.append(user)
    
    for user in users:
        session.add(user)
    
    await session.commit()
    print(f"创建了 {len(users)} 个测试用户")


async def create_test_manufacturers(session: AsyncSession):
    """创建测试厂商"""
    print("创建测试厂商...")
    
    # 获取厂商用户
    result = await session.execute(select(User).where(User.role == UserRole.MANUFACTURER))
    manufacturer_users = result.scalars().all()
    
    manufacturers = []
    for i, user in enumerate(manufacturer_users, 1):
        manufacturer = Manufacturer(
            company_name=f"专业制造厂商{i}",
            contact_person=user.full_name,
            email=user.email,
            phone=user.phone,
            address=f"制造基地{i}号",
            description=f"专业制造各类3D定制商品，拥有先进的3D打印设备和技术团队",
            supported_materials='["3D打印", "CNC加工", "激光切割"]',
            max_production_capacity=1000,
            production_lead_time=random.randint(5, 15),
            is_verified=True,
            is_active=True
        )
        manufacturers.append(manufacturer)
    
    for manufacturer in manufacturers:
        session.add(manufacturer)
    
    await session.commit()
    print(f"创建了 {len(manufacturers)} 个测试厂商")


async def create_test_products(session: AsyncSession):
    """创建测试商品"""
    print("创建测试商品...")
    
    # 获取分类
    result = await session.execute(select(Category))
    categories = result.scalars().all()
    
    if not categories:
        print("没有找到商品分类，先创建分类")
        return
    
    products = []
    product_names = [
        "个性化钥匙扣", "创意手机支架", "装饰摆件", "定制玩具", "实用日用品",
        "商务礼品", "创意礼品", "个性饰品", "定制模型", "艺术摆件"
    ]
    
    for i in range(20):
        category = random.choice(categories)
        product = Product(
            name=f"{random.choice(product_names)}-{i+1}",
            description=f"这是一个精美的{category.name}，可以根据您的需求进行个性化定制",
            price=random.randint(50, 500) + random.randint(0, 99) / 100,
            category_id=category.id,
            material="PLA环保材料",
            dimensions="10x10x5cm",
            weight=random.randint(50, 200),
            customizable_fields='{"text": {"type": "string", "max_length": 20, "required": false}, "color": {"type": "color", "options": ["red", "blue", "green", "yellow"], "required": true}, "size": {"type": "select", "options": ["small", "medium", "large"], "required": true}}',
            production_time_days=random.randint(3, 10),
            is_active=True,
            stock_quantity=random.randint(10, 100),
            min_order_quantity=1,
            tags=f'["{category.name}", "定制", "3D打印"]'
        )
        products.append(product)
    
    for product in products:
        session.add(product)
    
    await session.commit()
    print(f"创建了 {len(products)} 个测试商品")


async def create_test_product_images(session: AsyncSession):
    """创建测试商品图片"""
    print("创建测试商品图片...")
    
    # 获取商品
    result = await session.execute(select(Product))
    products = result.scalars().all()
    
    if not products:
        print("没有找到商品，先创建商品")
        return
    
    for product in products:
        # 每个商品创建2-5张图片
        for i in range(random.randint(2, 5)):
            image = ProductImage(
                product_id=product.id,
                uploader_id=1,  # 使用默认管理员用户
                original_filename=f"{product.name}_{i+1}.jpg",
                file_url=f"/static/images/placeholder_{product.name}_{i+1}.jpg",
                file_size=random.randint(50000, 200000),
                mime_type="image/jpeg",
                width=800,
                height=600,
                is_primary=(i == 0)
            )
            session.add(image)
    
    await session.commit()
    print("测试商品图片创建完成")


async def create_test_3d_models(session: AsyncSession):
    """创建测试3D模型"""
    print("创建测试3D模型...")
    
    # 获取商品
    result = await session.execute(select(Product))
    products = result.scalars().all()
    
    if not products:
        print("没有找到商品，先创建商品")
        return
    
    for product in products:
        # 获取商品的图片
        result = await session.execute(select(ProductImage).where(ProductImage.product_id == product.id))
        images = result.scalars().all()
        
        if not images:
            continue
            
        # 每个商品创建1-3个3D模型
        for i in range(random.randint(1, 3)):
            source_image = random.choice(images)
            model = Model3D(
                product_id=product.id,
                source_image_id=source_image.id,
                model_file_url=f"https://example.com/models/{product.name}_{i+1}.glb",
                file_format="glb",
                file_size=random.randint(1024, 10240 * 1024),
                original_filename=f"{product.name}_{i+1}.glb",
                storage_path=f"/models/{product.name}_{i+1}.glb",
                is_optimized=True,
                optimization_params='{"level": "high"}',
            )
            session.add(model)
    
    await session.commit()
    print("测试3D模型创建完成")


async def create_test_orders(session: AsyncSession):
    """创建测试订单"""
    print("创建测试订单...")
    
    # 获取用户
    result = await session.execute(select(User).where(User.role == UserRole.CUSTOMER))
    users = result.scalars().all()
    
    if not users:
        print("没有找到用户，先创建用户")
        return
    
    result = await session.execute(select(Product))
    products = result.scalars().all()
    
    if not users or not products:
        print("用户或商品不足，跳过订单创建")
        return
    
    orders = []
    statuses = [OrderStatus.PENDING, OrderStatus.PROCESSING, OrderStatus.SHIPPED, OrderStatus.DELIVERED]
    
    for i in range(50):
        user = random.choice(users)
        product = random.choice(products)
        quantity = random.randint(1, 5)
        
        # 随机生成订单日期（最近30天内）
        order_date = datetime.now() - timedelta(days=random.randint(0, 30))
        
        # 获取一个3D模型
        result = await session.execute(select(Model3D).where(Model3D.product_id == product.id))
        models = result.scalars().all()
        if not models:
            # 如果没有3D模型，创建一个新的
            model = Model3D(
                product_id=product.id,
                source_image_id=None,
                model_file_url=f"https://example.com/models/{product.name}_default.glb",
                file_format="glb",
                file_size=random.randint(1024, 10240 * 1024),
                is_optimized=True,
                optimization_params='{"level": "high"}',
            )
            session.add(model)
            await session.flush()
        else:
            model = random.choice(models)
        
        import uuid
        order_number = f"ORD-{uuid.uuid4().hex[:8].upper()}"
        order = Order(
            order_number=order_number,
            user_id=user.id,
            product_id=product.id,
            model_3d_id=model.id,
            quantity=quantity,
            unit_price=product.price,
            total_amount=product.price * quantity,
            shipping_address="测试收货地址",
            shipping_name=user.full_name or f"用户{user.id}",
            shipping_phone=user.phone or f"13800{user.id:05d}",
            custom_params=f'{{"text": "定制文字{i+1}", "color": "{random.choice(["red", "blue", "green"])}", "size": "{random.choice(["small", "medium", "large"])}"}}',
            status=random.choice(statuses),
            payment_status="pending",
            created_at=order_date
        )
        
        # 设置一些订单为已支付
        if random.random() > 0.3:
            order.payment_status = "paid"
            order.payment_method = random.choice(["alipay", "wechat", "card"])
            order.paid_at = order_date + timedelta(hours=random.randint(1, 24))
        
        orders.append(order)
    
    for order in orders:
        session.add(order)
    
    await session.flush()  # 确保订单被添加到会话
    await session.commit()
    print(f"创建了 {len(orders)} 个测试订单")


async def create_test_production_jobs(session: AsyncSession):
    """创建测试生产任务"""
    print("创建测试生产任务...")
    
    # 获取订单
    result = await session.execute(select(Order))
    orders = result.scalars().all()
    
    result = await session.execute(select(Manufacturer))
    manufacturers = result.scalars().all()
    
    if not orders or not manufacturers:
        print("订单或厂商不足，跳过生产任务创建")
        return
    
    jobs = []
    statuses = [JobStatus.PENDING, JobStatus.ASSIGNED, JobStatus.IN_PROGRESS, JobStatus.COMPLETED, JobStatus.FAILED]
    
    for order in orders:
        manufacturer = random.choice(manufacturers)
        
        # 获取订单对应的3D模型
        result = await session.execute(select(Model3D).where(Model3D.product_id == order.product_id))
        models = result.scalars().all()
        
        if not models:
            # 如果没有3D模型，创建默认模型
            model = Model3D(
                product_id=order.product_id,
                source_image_id=None,
                model_file_url=f"https://example.com/models/product_{order.product_id}_default.glb",
                file_format="glb",
                file_size=random.randint(1024, 10240 * 1024),
                original_filename=f"product_{order.product_id}_default.glb",
                storage_path=f"/models/product_{order.product_id}_default.glb",
                status="generated",
                is_optimized=True,
                optimization_params='{"level": "high"}',
            )
            session.add(model)
            await session.flush()
        else:
            model = random.choice(models)
        
        import uuid
        job_number = f"JOB-{uuid.uuid4().hex[:8].upper()}"
        job = ProductionJob(
            job_number=job_number,
            order_id=order.id,
            manufacturer_id=manufacturer.id,
            production_file_url=model.model_file_url,
            file_format=model.file_format,
            material_specification="PLA环保材料",
            quantity=order.quantity,
            priority=random.randint(1, 5),
            status=random.choice(statuses),
            estimated_completion_date=datetime.now() + timedelta(days=random.randint(1, 10))
        )
        
        # 设置一些任务为已完成
        if job.status == JobStatus.COMPLETED:
            job.completed_at = datetime.now() - timedelta(days=random.randint(1, 5))
            job.actual_completion_date = job.completed_at
        
        jobs.append(job)
    
    for job in jobs:
        session.add(job)
    
    await session.commit()
    print(f"创建了 {len(jobs)} 个测试生产任务")


async def generate_test_data():
    """生成所有测试数据"""
    print("开始生成测试数据...")
    
    # 创建数据库引擎
    engine = create_async_engine(DATABASE_URL, echo=True)
    async_session = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
    
    async with async_session() as session:
        try:
            await create_test_users(session)
            await create_test_manufacturers(session)
            await create_test_products(session)
            await create_test_product_images(session)
            await create_test_3d_models(session)
            await create_test_orders(session)
            await create_test_production_jobs(session)
            
            print("测试数据生成完成！")
            
        except Exception as e:
            print(f"测试数据生成失败: {e}")
            await session.rollback()
            raise
        finally:
            await session.close()
    
    await engine.dispose()


if __name__ == "__main__":
    print("3D定制商品生产平台 - 测试数据生成脚本")
    print("=" * 50)
    
    # 运行异步数据生成
    asyncio.run(generate_test_data())
    
    print("\n测试数据生成完成！")
    print("现在您可以使用测试账号登录系统：")
    print("普通用户: user1 / user1123")
    print("厂商用户: manufacturer1 / manufacturer1123")
    print("管理员: admin / admin123")